import sys
import os
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QHBoxLayout,
                             QWidget, QPushButton, QLabel, QLineEdit, QComboBox,
                             QProgressBar, QFileDialog, QMessageBox)
from PyQt5.QtCore import QThread, pyqtSignal
from PyQt5.QtGui import QFont
import serial
import serial.tools.list_ports
from intelhex import IntelHex
import time
from PyQt5.QtGui import QIcon

# YModem 协议实现
class YModem:
    SOH = 0x01
    STX = 0x02
    EOT = 0x04
    ACK = 0x06
    NAK = 0x15
    CAN = 0x18
    CRC16 = 0x43

    def __init__(self, serial_port):
        self.serial = serial_port
        self.block_number = 0

    def send_string(self, data):
        self.serial.write(data)
        return len(data)

    def send_block_end(self):
        # 发送数据块
        data = f"".encode('ascii')
        header = bytes([0x04])
        header += bytes([YModem.SOH])

        self.block_number = 0
        block_num = self.block_number % 256
        block_num_comp = 255 - block_num

        header += bytes([block_num, block_num_comp])

        # 填充数据块到固定大小

        data = data.ljust(1024, b'\x00')

        # 计算CRC
        crc = self.calc_crc16(data)
        packet = header + data + bytes([crc >> 8, crc & 0xFF])

        self.serial.write(packet)
        return packet

    def send_block(self, data):
        # 发送数据块
        # if len(data) <= 128:
        #     header = bytes([YModem.SOH])
        # else:
        header = bytes([YModem.STX])


        block_num = self.block_number % 256
        block_num_comp = 255 - block_num

        header += bytes([block_num, block_num_comp])

        # 填充数据块到固定大小
        if len(data) < 1024:
            data = data.ljust(1024, b'\x00')

        # 计算CRC
        crc = self.calc_crc16(data)
        packet = header + data + bytes([crc >> 8, crc & 0xFF])

        self.serial.write(packet)
        return packet

    def wait_ack_msg(self, timeout=1):
        # 等待ACK
        start_time = time.time()
        while time.time() - start_time < timeout:
            if self.serial.in_waiting > 0:
                return self.serial.read(1)
        return False

    def wait_ack(self, timeout=3):
        # 等待ACK
        start_time = time.time()
        while time.time() - start_time < timeout:
            if self.serial.in_waiting > 0:
                response = self.serial.read(1)
                if response == bytes([YModem.ACK]):
                    return True
                elif response == bytes([YModem.NAK]):
                    return False
                elif response == bytes([YModem.CAN]):
                    return None
        return False

    def send_file(self, file_data, progress_callback=None):
        # 发送文件
        file_size = len(file_data)
        file_name = "Demo.bin"

        # 发送文件名和大小
        file_info = f"{file_name}\0{file_size}".encode('ascii')
        # self.serial.write(bytes([YModem.SOH]))
        update_string = "update\r\n"
        self.send_string(f"{update_string}".encode('ascii'))

        # 发送文件信息块
        while True:
            if self.wait_ack_msg() == bytes([YModem.CRC16]):
                break
            # self.serial.write(bytes([YModem.CRC16]))

        self.block_number = 0
        self.send_block(file_info)
        if not self.wait_ack():
            return False

        # 发送文件数据
        sent_bytes = 0
        while sent_bytes < file_size:
            chunk_size = min(1024, file_size - sent_bytes)
            chunk = file_data[sent_bytes:sent_bytes + chunk_size]

            self.block_number += 1
            packet = self.send_block(chunk)

            if not self.wait_ack():
                return False

            sent_bytes += chunk_size
            if progress_callback:
                progress_callback(sent_bytes, file_size)

        # 发送结束标志
        # self.serial.write(bytes([YModem.EOT]))

        self.send_block_end()

        if not self.wait_ack():
            return False

        return True

    @staticmethod
    def calc_crc16(data):
        crc = 0
        for byte in data:
            crc ^= byte << 8
            for _ in range(8):
                if crc & 0x8000:
                    crc = (crc << 1) ^ 0x1021
                else:
                    crc <<= 1
                crc &= 0xFFFF
        return crc


class WorkerThread(QThread):
    progress_signal = pyqtSignal(int, int)
    result_signal = pyqtSignal(bool, str)

    def __init__(self, port_name, baud_rate, file_data):
        super().__init__()
        self.port_name = port_name
        self.baud_rate = baud_rate
        self.file_data = file_data
        self.running = True

    def run(self):
        try:
            with serial.Serial(self.port_name, self.baud_rate, timeout=1) as ser:
                ymodem = YModem(ser)
                success = ymodem.send_file(self.file_data, self.progress_callback)

                if success:
                    self.result_signal.emit(True, "文件传输成功!")
                else:
                    self.result_signal.emit(False, "文件传输失败!")
        except Exception as e:
            self.result_signal.emit(False, f"发生错误: {str(e)}")

    def progress_callback(self, sent, total):
        self.progress_signal.emit(sent, total)

    def stop(self):
        self.running = False




class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()

        self.setWindowIcon(QIcon("icon.png"))

        # 在窗口初始化时添加（如 __init__ 方法中）
        self.font = QFont()
        self.font.setPointSize(12)  # 默认通常为9，建议10-14之间
        self.font.setBold(True)
        self.setFont(self.font)  # 应用到整个窗口

        self.setWindowTitle("IAP TOOL")
        self.setGeometry(300, 300, 1000, 800)

        self.worker_thread = None
        self.file_data = None
        self.file_path = ""

        self.init_ui()
        self.refresh_ports()

    def init_ui(self):
        main_widget = QWidget()
        layout = QVBoxLayout()

        # 串口设置
        port_layout = QHBoxLayout()
        port_layout.setSpacing(40)
        port_layout.addWidget(QLabel("串口:"))
        self.port_combo = QComboBox()
        port_layout.addWidget(self.port_combo)

        refresh_btn = QPushButton("刷新")
        refresh_btn.clicked.connect(self.refresh_ports)
        port_layout.addWidget(refresh_btn)

        port_layout.addWidget(QLabel("波特率:"))
        self.baud_combo = QComboBox()
        self.baud_combo.addItems(["9600", "19200", "38400", "57600", "115200", "230400", "460800", "921600"])
        self.baud_combo.setCurrentText("115200")
        port_layout.addWidget(self.baud_combo)

        layout.addLayout(port_layout)

        # 文件选择
        file_layout = QHBoxLayout()
        file_layout.setSpacing(40)
        file_layout.addWidget(QLabel("固件文件:"))
        self.file_edit = QLineEdit()
        self.file_edit.setReadOnly(True)
        file_layout.addWidget(self.file_edit)

        browse_btn = QPushButton("浏览...")
        browse_btn.clicked.connect(self.browse_file)
        file_layout.addWidget(browse_btn)

        layout.addLayout(file_layout)

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        layout.addWidget(self.progress_bar)

        # 状态信息
        self.status_label = QLabel("准备就绪")
        layout.addWidget(self.status_label)

        # 操作按钮
        btn_layout = QHBoxLayout()
        btn_layout.setSpacing(40)
        self.start_btn = QPushButton("开始升级")
        self.start_btn.clicked.connect(self.start_upgrade)
        self.start_btn.setEnabled(False)
        btn_layout.addWidget(self.start_btn)

        self.cancel_btn = QPushButton("取消")
        self.cancel_btn.clicked.connect(self.cancel_upgrade)
        self.cancel_btn.setEnabled(False)
        btn_layout.addWidget(self.cancel_btn)

        layout.addLayout(btn_layout)

        main_widget.setLayout(layout)
        self.setCentralWidget(main_widget)

    def refresh_ports(self):
        self.port_combo.clear()
        ports = serial.tools.list_ports.comports()
        for port in ports:
            self.port_combo.addItem(port.device)

    def browse_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择固件文件", "", "固件文件 (*.bin *.hex);;所有文件 (*.*)")

        if file_path:
            self.file_path = file_path
            self.file_edit.setText(file_path)

            try:
                if file_path.lower().endswith('.hex'):
                    # 转换HEX为BIN
                    ih = IntelHex(file_path)
                    self.file_data = ih.tobinstr()
                    self.status_label.setText("HEX文件已加载并转换为BIN格式")
                else:
                    # 直接读取BIN文件
                    with open(file_path, 'rb') as f:
                        self.file_data = f.read()
                    self.status_label.setText("BIN文件已加载")

                self.start_btn.setEnabled(True)
            except Exception as e:
                QMessageBox.critical(self, "错误", f"加载文件失败: {str(e)}")
                self.file_data = None
                self.start_btn.setEnabled(False)

    def start_upgrade(self):
        if not self.file_data:
            QMessageBox.warning(self, "警告", "请先选择有效的固件文件")
            return

        if self.port_combo.count() == 0:
            QMessageBox.warning(self, "警告", "没有可用的串口")
            return

        port_name = self.port_combo.currentText()
        baud_rate = int(self.baud_combo.currentText())

        self.worker_thread = WorkerThread(port_name, baud_rate, self.file_data)
        self.worker_thread.progress_signal.connect(self.update_progress)
        self.worker_thread.result_signal.connect(self.upgrade_finished)
        self.worker_thread.start()

        self.start_btn.setEnabled(False)
        self.cancel_btn.setEnabled(True)
        self.status_label.setText("正在传输文件...")

    def cancel_upgrade(self):
        if self.worker_thread and self.worker_thread.isRunning():
            self.worker_thread.stop()
            self.worker_thread.wait()

        self.progress_bar.setValue(0)
        self.start_btn.setEnabled(True)
        self.cancel_btn.setEnabled(False)
        self.status_label.setText("传输已取消")

    def update_progress(self, sent, total):
        percent = int((sent / total) * 100)
        self.progress_bar.setValue(percent)
        self.status_label.setText(f"传输中: {sent}/{total} 字节 ({percent}%)")

    def upgrade_finished(self, success, message):
        self.progress_bar.setValue(100 if success else 0)
        self.status_label.setText(message)
        self.start_btn.setEnabled(True)
        self.cancel_btn.setEnabled(False)

        if success:
            QMessageBox.information(self, "成功", message)
        else:
            QMessageBox.warning(self, "失败", message)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())